"""
国密算法工具模块
实现AES、SM3、SM4算法的封装
"""
import base64
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes

class GMCrypto:
    """
    加密算法工具类
    """
    
    def __init__(self):
        """
        初始化加密算法工具
        """
        pass
    
    def aes_encrypt(self, data, key=None):
        """
        使用AES算法加密数据
        
        Args:
            data (str): 要加密的数据
            key (bytes, optional): AES密钥，如果未提供则自动生成
            
        Returns:
            tuple: (encrypted_data_base64, key_base64) 加密后的数据和密钥的Base64编码
        """
        # 如果没有提供密钥，则生成一个随机密钥
        if key is None:
            key = get_random_bytes(32)  # AES-256密钥长度为256位
        
        # 将数据转换为字节
        data_bytes = data.encode('utf-8')
        
        # 创建AES加密器（CBC模式）
        cipher = AES.new(key, AES.MODE_CBC)
        
        # 对数据进行PKCS7填充并加密
        padded_data = pad(data_bytes, AES.block_size)
        encrypted_data = cipher.encrypt(padded_data)
        
        # 将加密数据、IV和密钥转换为Base64编码
        encrypted_data_base64 = base64.b64encode(encrypted_data).decode('utf-8')
        iv_base64 = base64.b64encode(cipher.iv).decode('utf-8')
        key_base64 = base64.b64encode(key).decode('utf-8')
        
        return encrypted_data_base64, key_base64, iv_base64
    
    def aes_decrypt(self, encrypted_data_base64, key_base64, iv_base64):
        """
        使用AES算法解密数据
        
        Args:
            encrypted_data_base64 (str): Base64编码的加密数据
            key_base64 (str): Base64编码的密钥
            iv_base64 (str): Base64编码的初始化向量
            
        Returns:
            str: 解密后的数据
        """
        # 解码Base64数据
        encrypted_data = base64.b64decode(encrypted_data_base64)
        key = base64.b64decode(key_base64)
        iv = base64.b64decode(iv_base64)
        
        # 创建AES解密器（CBC模式）
        cipher = AES.new(key, AES.MODE_CBC, iv=iv)
        
        # 解密数据
        padded_data = cipher.decrypt(encrypted_data)
        
        # 移除PKCS7填充
        data_bytes = unpad(padded_data, AES.block_size)
        
        # 将解密后的数据转换为字符串
        data = data_bytes.decode('utf-8')
        
        return data
    
    def sm4_encrypt(self, data, key=None):
        """
        使用国密SM4算法加密数据
        
        Args:
            data (str): 要加密的数据
            key (bytes, optional): SM4密钥，如果未提供则自动生成
            
        Returns:
            tuple: (encrypted_data_base64, key_base64) 加密后的数据和密钥的Base64编码
        """
        # 如果没有提供密钥，则生成一个随机密钥
        if key is None:
            import os
            key = os.urandom(16)  # SM4密钥长度为128位
        
        # 将数据转换为字节
        data_bytes = data.encode('utf-8')
        
        # 创建SM4加密器
        from gmssl import sm4
        sm4_crypt = sm4.CryptSM4()
        sm4_crypt.set_key(key, sm4.SM4_ENCRYPT)
        
        # 对数据进行加密
        encrypted_data = sm4_crypt.crypt_ecb(data_bytes)
        
        # 将加密数据和密钥转换为Base64编码
        encrypted_data_base64 = base64.b64encode(encrypted_data).decode('utf-8')
        key_base64 = base64.b64encode(key).decode('utf-8')
        
        return encrypted_data_base64, key_base64
    
    def sm4_decrypt(self, encrypted_data_base64, key_base64):
        """
        使用国密SM4算法解密数据
        
        Args:
            encrypted_data_base64 (str): Base64编码的加密数据
            key_base64 (str): Base64编码的密钥
            
        Returns:
            str: 解密后的数据
        """
        # 解码Base64数据
        encrypted_data = base64.b64decode(encrypted_data_base64)
        key = base64.b64decode(key_base64)
        
        # 创建SM4解密器
        from gmssl import sm4
        sm4_crypt = sm4.CryptSM4()
        sm4_crypt.set_key(key, sm4.SM4_DECRYPT)
        
        # 解密数据
        decrypted_data = sm4_crypt.crypt_ecb(encrypted_data)
        
        # 将解密后的数据转换为字符串
        data = decrypted_data.decode('utf-8')
        
        return data